package pro.softzhang.algo.lc800;

import pro.softzhang.algo.common.TreeNode;
import pro.softzhang.algo.common.Utils;
import pro.softzhang.algo.flag.x.XClassic;

/**
 * 二叉树剪枝
 */
@XClassic
public class LC814_BinaryTreePruning {
    public static void main(String[] args) {
        TreeNode root = new TreeNode(1, new TreeNode(2, null, null), new TreeNode(3, null, null));
        Utils.printBinaryTree(root);
        root = new Solution().pruneTree(root);
        Utils.printBinaryTree(root);
    }

    static class Solution {
        public TreeNode pruneTree(TreeNode root) {
            if(root == null) {
                return root;
            }

            root.left = pruneTree(root.left);
            root.right = pruneTree(root.right);
            if(root.left == null && root.right == null && root.val == 0) {
                return null;
            }
            return root;
        }
    }
}